#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   write_range.py
@Time    :   2024/02/09 01:52:04
@Author  :   chakcy 
@Email   :   947105045@qq.com
@description   :   写入范围
'''

import pandas as pd  
from openpyxl import Workbook, load_workbook  
from openpyxl.utils import get_column_letter, column_index_from_string  
  
def write_to_excel_range(filename:str, sheet_name:str, data:list, range_string:str, include_header=True):  
    # 解析范围字符串，获取起始和结束的行列索引  
    start_col_row, end_col_row = range_string.split(':')  
    start_col, start_row = start_col_row[0], int(start_col_row[1:])  
    end_col, end_row = end_col_row[0], int(end_col_row[1:])  
    start_col_idx = column_index_from_string(start_col)  
    end_col_idx = column_index_from_string(end_col)  
  
    # 确保数据大小适合给定的范围  
    num_rows = end_row - start_row + 1  
    num_cols = end_col_idx - start_col_idx + 1  
    data = pd.DataFrame(data)  
    if data.shape != (num_rows, num_cols):  
        raise ValueError(  
            f"Data shape {data.shape} does not match the specified Excel range {range_string} ({num_rows} rows x {num_cols} columns).")  
  
    # 创建一个新的工作簿或加载现有的工作簿  
    try:  
        wb = load_workbook(filename)  
    except FileNotFoundError:  
        wb = Workbook()  
  
    # 选择或创建工作表  
    if sheet_name in wb.sheetnames:  
        ws = wb[sheet_name]  
    else:  
        ws = wb.create_sheet(sheet_name)  
  
    # 写入表头（如果include_header为True）  
    if include_header:  
        for c in range(num_cols):  
            col_letter = get_column_letter(start_col_idx + c)  
            ws[f'{col_letter}{start_row}'] = data.columns[c]  
        start_row += 1  # 更新起始行以跳过表头  
  
    # 将数据写入指定的范围（不包括表头）  
    for r in range(num_rows):  
        for c in range(num_cols):  
            col_letter = get_column_letter(start_col_idx + c)  
            ws[f'{col_letter}{start_row + r}'] = data.iloc[r, c]  
  
    # 保存工作簿  
    wb.save(filename)  
  
if __name__ == "__main__":
    # 示例数据  
    data = {  
        'A': [10, 11, 12],  
        'B': [20, 21, 22],  
        'C': [30, 31, 32]  
    }  
    
    # 调用函数，将数据写入"example.xlsx"的"Sheet1"工作表的A1:C3范围，附带表头  
    write_to_excel_range('example.xlsx', 'Sheet1', data, 'A1:C3', include_header=True)  
    
    # 如果不想附带表头，可以这样做：  
    # write_to_excel_range('example.xlsx', 'Sheet1', data, 'A1:C3', include_header=False)